package br.ufrn.avltree.impl;


import java.util.Stack;

import br.ufrn.binarytree.impl.BinarySearchTree;

public class AVLSearchTree<T extends Comparable<T>> extends BinarySearchTree<T>{
	private int height;
	
	public AVLSearchTree(){
		setKey(null);
		setLeft(null);
		setRight(null);
		height = 0;
	}
	
	public AVLSearchTree(T item){
		setKey(item);
		setLeft(new AVLSearchTree<T>());
		setRight(new AVLSearchTree<T>());
		this.adjustHeight();
	}
	
	public BinarySearchTree<T> insert(T item) throws Exception{
		Stack<AVLSearchTree<T>> elements = this.searchHelper(item);
		AVLSearchTree<T> node;
		AVLSearchTree<T> father;
		
		node = elements.pop();
		
		if(node.isEmpty()){
			node.insertHelper(item);
		}else{
			return this;
		}
		
		father = new AVLSearchTree<T>();
		
		while(!elements.isEmpty()){
			father = elements.pop();
			
			if(node.getKey().compareTo(father.getKey()) < 0){
				father.setLeft(node.balance());
			}else if(node.getKey().compareTo(father.getKey()) > 0){
				father.setRight(node.balance());
			}else{
				throw new Exception("Ilegal Tree Structure.");
			}
			
			father.adjustHeight();
			node = father;
		}
		if(!father.isEmpty()){
			father = father.balance();
		}else{
			father = node;
		}
		
		return father;
	}
	
	private void insertHelper(T item){
		this.setKey(item);
		this.setLeft(new AVLSearchTree<T>());
		this.setRight(new AVLSearchTree<T>());
		this.adjustHeight();
	}
	
	public AVLSearchTree<T> doRRRotation(){
		AVLSearchTree<T> oldRoot;
		AVLSearchTree<T> newRoot;
		AVLSearchTree<T> auxNode;
		
		oldRoot = this;
		newRoot = (AVLSearchTree<T>) oldRoot.getLeft();
		auxNode = (AVLSearchTree<T>) newRoot.getRight();
		
		oldRoot.setLeft(auxNode);
		newRoot.setRight(oldRoot);
		
		oldRoot.adjustHeight();
		newRoot.adjustHeight();
		
		return newRoot;
	}
	
	public AVLSearchTree<T> doLLRotation(){
		AVLSearchTree<T> oldRoot;
		AVLSearchTree<T> newRoot;
		AVLSearchTree<T> auxNode;
		
		oldRoot = this;
		newRoot = (AVLSearchTree<T>) oldRoot.getRight();
		auxNode = (AVLSearchTree<T>) newRoot.getLeft();
		
		oldRoot.setRight(auxNode);
		newRoot.setLeft(oldRoot);
		
		oldRoot.adjustHeight();
		newRoot.adjustHeight();
		
		return newRoot;
	}
	
	private AVLSearchTree<T> doLRRotation(){
		this.setLeft(((AVLSearchTree<T>) this.getLeft()).doLLRotation());
		return this.doRRRotation();
	}
	
	private AVLSearchTree<T> doRLRotation(){
		this.setRight(((AVLSearchTree<T>) this.getRight()).doRRRotation());
		return this.doLLRotation();
	}
	
	public void adjustHeight(){
		int h1 = ((AVLSearchTree<T>) this.getLeft()).getHeight();
		int h2 = ((AVLSearchTree<T>) this.getRight()).getHeight();
		
		this.setHeight((h1 > h2 ? h1 : h2) + 1);
	}
	
	public Stack<AVLSearchTree<T>> searchHelper(T item) throws Exception{
		Stack<AVLSearchTree<T>> elements = new Stack<AVLSearchTree<T>>();
		AVLSearchTree<T> auxNode;
		
		auxNode = this;
		
		elements.push(auxNode);
		
		while(!auxNode.isEmpty() && !auxNode.getKey().equals(item)){
			if(auxNode.getKey().compareTo(item) > 0){
				auxNode = (AVLSearchTree<T>) auxNode.getLeft();
			}else if(auxNode.getKey().compareTo(item) < 0){
				auxNode = (AVLSearchTree<T>) auxNode.getRight();
			}
			elements.push(auxNode);
		}
		
		return elements;
	}
	
	public AVLSearchTree<T> balance() throws Exception{
		if(this.getBalanceFactor() >= 2){
			if(((AVLSearchTree<T>) this.getLeft()).getBalanceFactor() >= 1){
				return this.doRRRotation();
			}else if(((AVLSearchTree<T>) this.getLeft()).getBalanceFactor() <= -1){
				return this.doLRRotation();
			}else{
				throw new Exception("Invalid balance.");
			}
		}else if(this.getBalanceFactor() <= -2){
			if(((AVLSearchTree<T>) this.getRight()).getBalanceFactor() <= -1){
				return this.doLLRotation();
			}else if(((AVLSearchTree<T>) this.getRight()).getBalanceFactor() >= 1){
				return this.doRLRotation();
			}else{
				throw new Exception("Invalid balance.");
			}	
		}else{
			return this;
		}
	}
	
	private int getBalanceFactor(){
		int h1 = ((AVLSearchTree<T>) this.getLeft()).getHeight();
		int h2 = ((AVLSearchTree<T>) this.getRight()).getHeight();
		
		return h1 - h2;
	}

	public void setHeight(int height) {
		this.height = height;
	}

	public int getHeight() {
		return height;
	}
}